/* Actual UDP implementation */

/* Copyright (C) 2009-2011 David Zanetti
 *
 *  This program is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; version 2 of the License.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License along
 *  with this program; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
 */

#include <stdio.h>
#include <avr/io.h>
#include "global.h"
#include <string.h>
#include "checksum.h"
#include "encx24j600.h"
#include <util/delay.h>
#include "serial.h"
#include <avr/pgmspace.h>
#include "udp.h"
#include "ntp.h"
#include "ip.h"
#include "dhcpc.h"
#include "mdns.h"
#include "tftp.h"

typedef struct {
		uint16_t source_port;
		uint16_t dest_port;
		uint16_t length;
		uint16_t checksum;
} udp_header_t;

/* compute header CRC, which includes psuedoip header */
/* note: length here is PAYLOAD length, not length of header */
uint16_t udp_psuedocrc(uint16_t length, uint8_t *src, uint8_t *dst,
		uint8_t *pkt);

uint8_t udp_process(uint8_t slot, uint8_t *src, uint8_t *dst) {
	udp_header_t pkt;
	uint16_t crc;
	uint16_t len;
	uint8_t ret = 0;
	uint16_t s;

	/* retrieve the UDP header */
	if (encx24j600_rx(&pkt, sizeof(udp_header_t)) != sizeof(udp_header_t)) {
		serial0_tx_PGM(PSTR("udp: short frame, discarding\r\n"));
		return E_NOTHING;
	}

	s = pkt.source_port;

	/* since we use length quite a lot, but it's in network order, remap it
	 here */
	len = swap16(pkt.length);

	/* UDP CRC consists of the IP Pseudo Header, UDP header, and all
	 payload. We compute the psuedoheader+udpheader locally since it's only
	 a few bytes and too short to write somewhere useful on the NIC */

	/* psuedoheader and main udp header */
	crc = udp_psuedocrc(len, src, dst, (uint8_t *) &pkt);
	/* payload only .. per 14.2 copes with odd-length payload */
	crc = encx24j600_rx_crc_unread(len - sizeof(udp_header_t), swap16(crc));

#ifdef DEBUG_UDP
	serial0_tx_PGM(PSTR("udp: "));
	serial0_tx_dec(swap16(pkt.source_port));
	serial0_tx_PGM(PSTR("->"));
	serial0_tx_dec(swap16(pkt.dest_port));
	serial0_tx_PGM(PSTR(" ["));
	serial0_tx_dec(swap16(pkt.length));
	serial0_tx_PGM(PSTR(" bytes]"));
	if (crc == 0x0) {
		serial0_tx_PGM(PSTR(" crc good\r\n"));
	} else {
		serial0_tx_PGM(PSTR(" crc bad\r\n"));
	}
#endif

	/* this only applies if we had a checksum in the inbound packet.
	 udp allows zero checksum to mean "don't care" */
	if (pkt.checksum && crc != 0) {
		serial0_tx_PGM(PSTR("udp: discarding corrupt packet\r\n"));
		return E_NOTHING;
	}

	switch (pkt.dest_port) {
		case swap16(123):
			/* Network Time Protocol */
#ifdef DEBUG_NTP
			serial0_tx_PGM(PSTR("udp: recieved ntp request\r\n"));
#endif
			ret = ntp_process(slot);
			break;
		case swap16(68): /* BOOTP/DHCP Client */
#ifdef DEBUG_DHCPC
			serial0_tx_PGM(PSTR("udp: received dhcp client packet\r\n"));
#endif
			ret = dhcpc_process(slot);
			break;
#ifdef MDNS_H
		case swap16(5353): /* mDNS, DNS-SD by mDNS */
			/* discard non-mDNS multicast address target */
			if (dst[0] == 224 && dst[1] == 0 && dst[2] == 0 && dst[3] == 251) {
				serial0_tx_PGM(PSTR("udp: received validish mDNS packet\r\n"));
				ret = mdns_process(slot);
			}
			break;
#endif
#ifdef TFTP_H
		case swap16(69): /* TFTP */
			ret = tftp_process(slot,src,len-sizeof(udp_header_t));
			break;
#endif
	}

	if (ret) {
		udp_header_t tx;
		uint16_t newlen;

		newlen = (encx24j600_slot_len(slot) + sizeof(udp_header_t));

		/* we have a packet to transmit, populate our header */
		tx.dest_port = pkt.source_port;
		tx.source_port = pkt.dest_port;
		tx.checksum = 0;
		tx.length = swap16(newlen);

		if (ret == E_REPLY_BROAD) {
			uint8_t broad[4] = { 0xff, 0xff, 0xff, 0xff };
			/* need a different calculation for checksum */
			crc = udp_psuedocrc(newlen,ip_addr,broad,(uint8_t *)&tx);
		} else {
			/* compute checksum of header and psuedoheader */
			crc = udp_psuedocrc(newlen,ip_addr,src,(uint8_t *)&tx);
		}

		/* and then CRC the payload */
		crc = encx24j600_tx_crc_slot(slot,newlen-sizeof(udp_header_t),swap16(crc));

		/* now insert that CRC into our header */
		tx.checksum = crc;

#ifdef DEBUG_UDP
		serial0_tx_PGM(PSTR("udp: "));
		serial0_tx_dec(swap16(tx.source_port));
		serial0_tx_PGM(PSTR("->"));
		serial0_tx_dec(swap16(tx.dest_port));
		serial0_tx_PGM(PSTR(" ["));
		serial0_tx_dec(swap16(tx.length));
		serial0_tx_PGM(PSTR(" bytes, crc "));
		serial0_tx_hex(crc & 0xff);
		serial0_tx_hex(crc >> 8);
		serial0_tx_cout(']');
		serial0_tx_cr();
#endif

		/* prepend our header, and pass it up the stack */
		if (!encx24j600_tx_prepend(slot,(uint8_t *)&tx,sizeof(udp_header_t))) {
			serial0_tx_PGM(PSTR("udp: failed to prepend our header\r\n"));
			return E_NOTHING;
		}
		return ret; /* return the hint we got */
	}

	return E_NOTHING;
}

/* return the initial sums of the psuedoheader, no inversion */
uint16_t udp_psuedocrc(uint16_t length, uint8_t *src, uint8_t *dst,
		uint8_t *pkt) {
	uint32_t sum = 0;
	uint8_t n = sizeof(udp_header_t);

	/* all this casting is needed to avoid problems with avr-gcc generating 32-bit math */

	/* psuedoheader */
	sum += (uint32_t) ((uint32_t) src[0] << 8 | (uint32_t) src[1]);
	sum += (uint32_t) ((uint32_t) src[2] << 8 | (uint32_t) src[3]);

	sum += (uint32_t) ((uint32_t) dst[0] << 8 | (uint32_t) dst[1]);
	sum += (uint32_t) ((uint32_t) dst[2] << 8 | (uint32_t) dst[3]);

	sum += ((uint32_t) 17); /* UDP protocol */
	sum += ((uint32_t) length);

	/* header */
	while (n > 1) {
		/* all this casting is needed to force avr-gcc to actually use 32 bits */
		sum += (uint32_t)(((uint32_t) * pkt << 8) | ((uint32_t) * (pkt + 1)));
		pkt += 2;
		n -= 2;
	}

	/* and the final carries */
	if (sum >> 16) {
		sum = (sum & 0xffff) + (sum >> 16);
	}

	sum = ~sum;

	return ((uint16_t)(sum));

}

/* Send an arbitrary UDP packet to the given IP address. If you don't care what the source port is, randomly generate one. */
uint8_t udp_send(uint8_t *dst_ip, uint16_t port, uint16_t s_port, uint8_t *pkt,
		uint16_t length) {
	udp_header_t tx; /* our header to transmit */
	uint16_t crc; /* CRC for our packet */
	uint8_t slot; /* our TX slot */

	/* FIXME: we need to check if link is up really.. */

	slot = encx24j600_tx_freeslot();

	if (slot == 0xff) {
		serial0_tx_PGM(PSTR("udp: send failed; no free TX slots\r\n"));
		return 0;
	}

	/* push in the payload first */
	if (!encx24j600_tx_prepend(slot,(uint8_t *)pkt,length)) {
		serial0_tx_PGM(PSTR("udp: packet too big, discarded\r\n"));
		return 0;
	}

	/* build our UDP header for this packet */
	tx.dest_port = swap16(port);
	tx.source_port = swap16(s_port);
	tx.length = swap16((length+sizeof(udp_header_t)));
	tx.checksum = 0;

	/* compute checksum of header and psuedoheader */
	crc = udp_psuedocrc(length+sizeof(udp_header_t),ip_addr,dst_ip,(uint8_t *)&tx);

	/* and then CRC the payload */
	crc = encx24j600_tx_crc_slot(slot,length,swap16(crc));

	/* now insert that CRC into our header */
	tx.checksum = crc;

#ifdef DEBUG_UDP
	serial0_tx_PGM(PSTR("udp: "));
	serial0_tx_dec(swap16(tx.source_port));
	serial0_tx_PGM(PSTR("->"));
	serial0_tx_dec(swap16(tx.dest_port));
	serial0_tx_PGM(PSTR(" ["));
	serial0_tx_dec(swap16(tx.length));
	serial0_tx_PGM(PSTR(" bytes, crc "));
	serial0_tx_hex(crc & 0xff);
	serial0_tx_hex(crc >> 8);
	serial0_tx_cout(']');
	serial0_tx_cr();
#endif

	/* append our header and offer it to the next layer */
	if (encx24j600_tx_prepend(slot,(uint8_t *)&tx,sizeof(udp_header_t))) {
		/* insert our calculated CRC into the UDP packet */
		if (ip_send(dst_ip,slot,17) == E_NOTHING) { /* get IP layer to send us some protocol 17! */
			serial0_tx_PGM(PSTR("udp: failed to send\r\n"));
			encx24j600_slot_done(slot); /* so we don't leak the slot */
		} else {
			/* this should only happen when ARP is doing delayed TX */
#ifdef DEUBG_UDP
			serial0_tx_PGM(PSTR("udp: delayed TX from ARP resolution\r\n"));
#endif
			return 1;
		}
		return 1;
	} else {
		encx24j600_slot_done(slot);
		serial0_tx_PGM(PSTR("udp: failed to prepend header, giving up send\r\n"));
	}
	return 0;
}
